import joblib
import pandas as pd
import streamlit as st
from models import health_model_create, salary_model_create

st.title("机器学习实训")

# 创建一个侧边栏
st.sidebar.title('选择任务')


# 定义页面跳转的函数
def show_content(page):
    if page == '中风预测':
        df = pd.DataFrame(columns=joblib.load('./models/health_model/health_columns.pkl'))

        st.subheader('中风预测')
        gender = st.radio('性别', ('男', '女', '其他'))
        if gender == '男':
            st.write("您是男性。")
            df['Male'] = [1]
            df['Female'] = [0]
            df['Other'] = [0]
        elif gender == '女':
            st.write("您是女性。")
            df['Female'] = [1]
            df['Male'] = [0]
            df['Other'] = [0]
        else:
            st.write("其他。")
            df['Other'] = 1
            df['Male'] = 0
            df['Female'] = 0

        age = st.slider('年龄', 1, 100, 18)
        st.write("您的年龄是：", age)
        df['age'] = age

        hypertension = st.checkbox('是否有高血压')
        if hypertension:
            st.write("您有高血压。")
            df['hypertension'] = 1
        else:
            st.write("您没有高血压。")
            df['hypertension'] = 0

        heart_disease = st.checkbox('是否有心脏病')
        if heart_disease:
            st.write("您有心脏病。")
            df['heart_disease'] = 1
        else:
            st.write("您没有心脏病。")
            df['heart_disease'] = 0

        ever_married = st.checkbox('是否结婚')
        if ever_married:
            st.write("您已结婚。")
            df['Yes'] = 1
            df['No'] = 0
        else:
            st.write("您未结婚。")
            df['No'] = 1
            df['Yes'] = 0

        work_type = st.selectbox('工作类型',
                                 ('童工', '公务员', '未工作', '自营职业', '私营职业'))
        st.write("您的工作类型是：", work_type)
        if work_type == '童工':
            df['children'] = 1
        elif work_type == '公务员':
            df['Govt_job'] = 1
        elif work_type == '未工作':
            df['Never_worked'] = 1
        elif work_type == '自营职业':
            df['Self_employed'] = 1
        elif work_type == '私营职业':
            df['Private'] = 1

        residence_type = st.radio('居住类型', ('城市', '农村'))
        if residence_type == '城市':
            st.write("您居住在城市。")
            df['Urban'] = 1
            df['Rural'] = 0
        else:
            st.write("您居住在农村。")
            df['Rural'] = 1
            df['Urban'] = 0

        avg_glucose_level = st.number_input("请输入血液中的平均葡萄糖水平：")
        st.write("您的平均葡萄糖水平是：", avg_glucose_level)
        df['avg_glucose_level'] = avg_glucose_level

        bmi = st.number_input("请输入您的BMI：")
        st.write("您的BMI是：", bmi)
        df['bmi'] = bmi

        smoke = st.radio('是否吸烟',
                         ('以前吸烟过', '从未吸烟', '吸烟'))
        st.write("您的吸烟情况是：", smoke)
        if smoke == '以前吸烟过':
            df['formerly smoked'] = 1
        elif smoke == '从未吸烟':
            df['never smoked'] = 1
        elif smoke == '吸烟':
            df['smokes'] = 1

        model = None
        model_selection = st.selectbox('选择模型',
                                       ('逻辑回归', '决策树算法', 'KNN算法', '贝叶斯算法', '支持向量机', '随机森林'))
        if model_selection == '逻辑回归':
            model = joblib.load('./models/health_model/logistic_regression.pkl')
        elif model_selection == '决策树算法':
            model = joblib.load('./models/health_model/decision_tree.pkl')
        elif model_selection == 'KNN算法':
            model = joblib.load('./models/health_model/KNN.pkl')
        elif model_selection == '贝叶斯算法':
            model = joblib.load('./models/health_model/Bayesian.pkl')
        elif model_selection == '支持向量机':
            model = joblib.load('./models/health_model/SVM.pkl')
        elif model_selection == '随机森林':
            model = joblib.load('./models/health_model/random_forest.pkl')

        print(df)
        df = df.fillna(0)
        print(df)
        y_pred = model.predict(df)
        print(y_pred)

        if st.button('预测'):
            if y_pred == 1:
                st.markdown('<span style="color: red;">您的风险级别为高。</span>', unsafe_allow_html=True)
            else:
                st.markdown('<span style="color: green;">您的风险级别为低。</span>', unsafe_allow_html=True)
    elif page == '薪资预测':
        st.subheader('薪资预测')
        df = pd.DataFrame(columns=joblib.load('./models/salary_model/salary_columns.pkl'))
        df['fnlwgt'] = 189793

        sex = st.radio('性别', ('男', '女'))
        if sex == '男':
            st.write("您是男性。")
            df['sex_ Male'] = [1]
            df['sex_ Female'] = [0]
        elif sex == '女':
            st.write("您是女性。")
            df['sex_ Female'] = [1]
            df['sex_ Male'] = [0]

        age = st.slider('年龄', 1, 100, 18)
        st.write("您的年龄是：", age)
        df['age'] = age

        workclass = st.selectbox('工人阶级：',
                                 ('私人', '自雇非公司', '自雇公司', '联邦政府', '地方政府', '州政府', '无薪',
                                  '从未工作过'))
        if workclass == '私人':
            df['workclass_ Private'] = 1
        elif workclass == '自雇非公司':
            df['workclass_ Self-emp-not-inc'] = 1
        elif workclass == '自雇公司':
            df['workclass_ Self-emp-inc'] = 1
        elif workclass == '联邦政府':
            df['workclass_ Federal-gov'] = 1
        elif workclass == '地方政府':
            df['workclass_ Local-gov'] = 1
        elif workclass == '州政府':
            df['workclass_ State-gov'] = 1
        elif workclass == '无薪':
            df['workclass_ Without-pay'] = 1
        elif workclass == '从未工作过':
            df['workclass_ Never-worked'] = 1

        education = st.selectbox('学历：',
                                 ("Bachelors", "Some-college", "11th", "HS-grad", "Prof-school", "Assoc-acdm",
                                  "Assoc-voc", "9th", "7th-8th", "12th", "Masters", "1st-4th", "10th", "Doctorate",
                                  "5th-6th", "Preschool"))
        st.write("您的学历是：", education)
        df['education_ ' + education] = 1

        education_num = st.slider('受教育年限', 0, 30, 9)
        st.write("您的受教育年限是：", education_num)
        df['education-num'] = education_num

        marital_status = st.selectbox('婚姻状况',
                                      ('已婚配偶', '离异', '未婚', '分居', '丧偶', '已婚配偶缺席', '未知'))
        if marital_status == '已婚配偶':
            df['marital-status_ Married-civ-spouse'] = 1
        elif marital_status == '离异':
            df['marital-status_ Divorced'] = 1
        elif marital_status == '未婚':
            df['marital-status_ Never-married'] = 1
        elif marital_status == '分居':
            df['marital-status_ Separated'] = 1
        elif marital_status == '丧偶':
            df['marital-status_ Widowed'] = 1
        elif marital_status == '已婚配偶缺席':
            df['marital-status_ Married-spouse-absent'] = 1
        elif marital_status == '未知':
            df['marital-status_ Married-AF-spouse'] = 1

        relationship = st.selectbox('关系',
                                    ('妻子', '自己的孩子', '丈夫', '不属于家族', "其他关系", '未婚'))
        st.write("与您的关系是：", relationship)
        if relationship == '妻子':
            df['relationship_ Wife'] = 1
        elif relationship == '自己的孩子':
            df['relationship_ Own-child'] = 1
        elif relationship == '丈夫':
            df['relationship_ Husband'] = 1
        elif relationship == '不属于家族':
            df['relationship_ Not-in-family'] = 1
        elif relationship == '其他关系':
            df['relationship_ Other-relative'] = 1
        elif relationship == '未婚':
            df['relationship_ Unmarried'] = 1

        race = st.selectbox('种族',
                            ('White', 'Asian-Pac-Islander', 'Amer-Indian-Eskimo', 'Other', 'Black'))
        st.write("您的种族是：", race)
        df['race_ ' + race] = 1

        capital_gain = st.number_input("请输入公司收益：")
        st.write("您公司收入为：", capital_gain)
        df['capital-gain'] = capital_gain

        capital_loss = st.number_input("请输入公司损失：")
        st.write("您公司损失为：", capital_loss)
        df['capital-loss'] = capital_loss

        hours_per_week = st.number_input("请输入您每周工作小时数：")
        st.write("您每周工作小时数为：", hours_per_week)
        df['hours-per-week'] = hours_per_week

        native_country = st.selectbox('您的国家',
                                      ('United-States', 'Cambodia', 'England', 'Puerto-Rico', 'Canada', 'Germany',
                                       'Outlying-US(Guam-USVI-etc)', 'India', 'Japan', 'Greece', 'South', 'China',
                                       'Cuba', 'Iran', 'Honduras', 'Philippines', 'Italy', 'Poland', 'Jamaica',
                                       'Vietnam', 'Mexico', 'Portugal', 'Ireland', 'France', 'Dominican-Republic',
                                       'Laos', 'Ecuador', 'Taiwan', 'Haiti', 'Columbia', 'Hungary', 'Guatemala',
                                       'Nicaragua', 'Scotland', 'Thailand', 'Yugoslavia', 'El-Salvador',
                                       'Trinadad&Tobago', 'Peru', 'Hong', 'Holand-Netherlands'))
        st.write("您的国家是：", native_country)
        df['native-country_ ' + native_country] = 1

        model = None
        model_selection = st.selectbox('选择模型',
                                       ('逻辑回归', '决策树算法', 'KNN算法', '贝叶斯算法', '支持向量机', '随机森林'))
        if model_selection == '逻辑回归':
            model = joblib.load('./models/salary_model/logistic_regression.pkl')
        elif model_selection == '决策树算法':
            model = joblib.load('./models/salary_model/decision_tree.pkl')
        elif model_selection == 'KNN算法':
            model = joblib.load('./models/salary_model/KNN.pkl')
        elif model_selection == '贝叶斯算法':
            model = joblib.load('./models/salary_model/Bayesian.pkl')
        elif model_selection == '支持向量机':
            model = joblib.load('./models/salary_model/SVM.pkl')
        elif model_selection == '随机森林':
            model = joblib.load('./models/salary_model/random_forest.pkl')

        print(df)
        df = df.fillna(0)
        print(df)
        y_pred = model.predict(df)
        print(y_pred)

        if st.button('预测'):
            if y_pred == 1:
                st.markdown('<span style="color: blue;">您的薪资>50k。</span>', unsafe_allow_html=True)
            else:
                st.markdown('<span style="color: green;">您的薪资<=50k。</span>', unsafe_allow_html=True)
    elif page == '模型训练':
        st.subheader('模型训练')
        file_up = st.file_uploader("上传数据集", type=['csv'])
        if file_up:
            st.write(f"文件已上传: {file_up.name}")
            df = pd.read_csv(file_up)
            if st.button('开始训练'):
                if 'gender' in df.columns:
                    with st.spinner('正在处理...'):
                        health_model_create.train_and_evaluate_models(df)
                        st.write("模型训练完成。")
                elif 'sex' in df.columns:
                    with st.spinner('正在处理...'):
                        salary_model_create.train_and_evaluate_salary_models(df)
                        st.write("模型训练完成。")
                else:
                    st.write("请上传符合的数据集。")
        else:
            st.write("没有文件被上传。")


# 让用户在侧边栏选择页面
page = st.sidebar.selectbox(
    '选择一个任务',
    ('中风预测', '薪资预测', '模型训练')
)

# 根据用户选择显示相应内容
show_content(page)
